{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Copyright (c) Facebook, Inc. and its affiliates.\n",
    "# Licensed under the Apache License, Version 2.0 (the \"License\");\n",
    "# you may not use this file except in compliance with the License.\n",
    "# You may obtain a copy of the License at\n",
    "#\n",
    "#    http://www.apache.org/licenses/LICENSE-2.0\n",
    "#\n",
    "# Unless required by applicable law or agreed to in writing, software\n",
    "# distributed under the License is distributed on an \"AS IS\" BASIS,\n",
    "# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n",
    "# See the License for the specific language governing permissions and\n",
    "# limitations under the License."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# KubeFlow Pipelines :  Pytorch Cifar10 Image classification\n",
    "\n",
    "This notebook shows PyTorch CIFAR10 end-to-end  classification example using Kubeflow Pipelines. \n",
    "\n",
    "An example notebook that demonstrates how to:\n",
    "\n",
    "* Get different tasks needed for the pipeline\n",
    "* Create a Kubeflow pipeline\n",
    "* Include Pytorch KFP components to preprocess, train, visualize and deploy the model in the pipeline\n",
    "* Submit a job for execution\n",
    "* Query(prediction and explain) the final deployed model\n",
    "* Interpretation of the model using the Captum Insights\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## import the necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [],
   "source": [
    "import kfp\n",
    "import json\n",
    "import os\n",
    "from kfp.onprem import use_k8s_secret\n",
    "from kfp import components\n",
    "from kfp.components import load_component_from_file, load_component_from_url\n",
    "from kfp import dsl\n",
    "from kfp import compiler\n",
    "\n",
    "import numpy as np\n",
    "import logging\n",
    "\n",
    "\n",
    "kfp.__version__"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Enter your gateway and the auth token\n",
    "[Use this extension on chrome to get token]( https://chrome.google.com/webstore/detail/editthiscookie/fngmhnnpilhplaeedifhccceomclgfbg?hl=en)\n",
    "\n",
    "![image.png](./image.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Update values for the ingress gateway and auth session"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "INGRESS_GATEWAY='http://istio-ingressgateway.istio-system.svc.cluster.local'\n",
    "AUTH=\"<enter your auth token>\"\n",
    "NAMESPACE=\"kubeflow-user-example-com\"\n",
    "COOKIE=\"authservice_session=\"+AUTH\n",
    "EXPERIMENT=\"Default\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set  the Log bucket and Tensorboard Image"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "MINIO_ENDPOINT=\"http://minio-service.kubeflow:9000\"\n",
    "LOG_BUCKET=\"mlpipeline\"\n",
    "TENSORBOARD_IMAGE=\"public.ecr.aws/pytorch-samples/tboard:latest\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set the client and create the experiment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "client = kfp.Client(host=INGRESS_GATEWAY+\"/pipeline\", cookies=COOKIE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<a href=\"http://istio-ingressgateway.istio-system.svc.cluster.local/pipeline/#/experiments/details/aac96a63-616e-4d88-9334-6ca8df2bb956\" target=\"_blank\" >Experiment details</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "data": {
      "text/plain": [
       "{'created_at': datetime.datetime(2021, 4, 22, 8, 44, 39, tzinfo=tzlocal()),\n",
       " 'description': None,\n",
       " 'id': 'aac96a63-616e-4d88-9334-6ca8df2bb956',\n",
       " 'name': 'Default',\n",
       " 'resource_references': [{'key': {'id': 'kubeflow-user-example-com',\n",
       "                                  'type': 'NAMESPACE'},\n",
       "                          'name': None,\n",
       "                          'relationship': 'OWNER'}],\n",
       " 'storage_state': 'STORAGESTATE_AVAILABLE'}"
      ]
     },
     "execution_count": 17,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "client.create_experiment(EXPERIMENT)\n",
    "experiments = client.list_experiments(namespace=NAMESPACE)\n",
    "my_experiment = experiments.experiments[0]\n",
    "my_experiment"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Set  the Inference parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "DEPLOY_NAME=\"torchserve\"\n",
    "MODEL_NAME=\"cifar10\"\n",
    "ISVC_NAME=DEPLOY_NAME+\".\"+NAMESPACE+\".\"+\"example.com\"\n",
    "INPUT_REQUEST=\"https://raw.githubusercontent.com/kubeflow/pipelines/master/samples/contrib/pytorch-samples/cifar10/input.json\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Load the the components yaml files for setting up the components"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "prepare_tensorboard_op = load_component_from_file(\"./common/tensorboard/component.yaml\")\n",
    "prep_op = components.load_component_from_file(\n",
    "    \"./cifar10/yaml/pre_process/component.yaml\"\n",
    ")\n",
    "train_op = components.load_component_from_file(\n",
    "    \"./cifar10/yaml/train/component.yaml\"\n",
    ")\n",
    "deploy_op = load_component_from_file(\"./common/deploy/component.yaml\")\n",
    "pred_op = load_component_from_file(\"./common/prediction/component.yaml\")\n",
    "minio_op = components.load_component_from_file(\n",
    "    \"./common/minio/component.yaml\"\n",
    ")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Define the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "@dsl.pipeline(\n",
    "    name=\"Training Cifar10 pipeline\", description=\"Cifar 10 dataset pipeline\"\n",
    ")\n",
    "def pytorch_cifar10( # pylint: disable=too-many-arguments\n",
    "    minio_endpoint=MINIO_ENDPOINT,\n",
    "    log_bucket=LOG_BUCKET,\n",
    "    log_dir=f\"tensorboard/logs/{dsl.RUN_ID_PLACEHOLDER}\",\n",
    "    mar_path=f\"mar/{dsl.RUN_ID_PLACEHOLDER}/model-store\",\n",
    "    config_prop_path=f\"mar/{dsl.RUN_ID_PLACEHOLDER}/config\",\n",
    "    model_uri=f\"s3://mlpipeline/mar/{dsl.RUN_ID_PLACEHOLDER}\",\n",
    "    tf_image=TENSORBOARD_IMAGE,\n",
    "    deploy=DEPLOY_NAME,\n",
    "    isvc_name=ISVC_NAME,\n",
    "    model=MODEL_NAME,\n",
    "    namespace=NAMESPACE,\n",
    "    confusion_matrix_log_dir=f\"confusion_matrix/{dsl.RUN_ID_PLACEHOLDER}/\",\n",
    "    checkpoint_dir=\"checkpoint_dir/cifar10\",\n",
    "    input_req=INPUT_REQUEST,\n",
    "    cookie=COOKIE,\n",
    "    ingress_gateway=INGRESS_GATEWAY,\n",
    "):\n",
    "\n",
    "    def sleep_op(seconds):\n",
    "        \"\"\"Sleep for a while.\"\"\"\n",
    "        return dsl.ContainerOp(\n",
    "            name=\"Sleep \" + str(seconds) + \" seconds\",\n",
    "            image=\"python:alpine3.6\",\n",
    "            command=[\"sh\", \"-c\"],\n",
    "            arguments=[\n",
    "                'python -c \"import time; time.sleep($0)\"',\n",
    "                str(seconds)\n",
    "            ],\n",
    "        )\n",
    "\n",
    "    \"\"\"This method defines the pipeline tasks and operations\"\"\"\n",
    "    pod_template_spec = json.dumps({\n",
    "        \"spec\": {\n",
    "            \"containers\": [{\n",
    "                \"env\": [\n",
    "                    {\n",
    "                        \"name\": \"AWS_ACCESS_KEY_ID\",\n",
    "                        \"valueFrom\": {\n",
    "                            \"secretKeyRef\": {\n",
    "                                \"name\": \"mlpipeline-minio-artifact\",\n",
    "                                \"key\": \"accesskey\",\n",
    "                            }\n",
    "                        },\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"AWS_SECRET_ACCESS_KEY\",\n",
    "                        \"valueFrom\": {\n",
    "                            \"secretKeyRef\": {\n",
    "                                \"name\": \"mlpipeline-minio-artifact\",\n",
    "                                \"key\": \"secretkey\",\n",
    "                            }\n",
    "                        },\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"AWS_REGION\",\n",
    "                        \"value\": \"minio\"\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"S3_ENDPOINT\",\n",
    "                        \"value\": f\"{minio_endpoint}\",\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"S3_USE_HTTPS\",\n",
    "                        \"value\": \"0\"\n",
    "                    },\n",
    "                    {\n",
    "                        \"name\": \"S3_VERIFY_SSL\",\n",
    "                        \"value\": \"0\"\n",
    "                    },\n",
    "                ]\n",
    "            }]\n",
    "        }\n",
    "    })\n",
    "\n",
    "    prepare_tb_task = prepare_tensorboard_op(\n",
    "        log_dir_uri=f\"s3://{log_bucket}/{log_dir}\",\n",
    "        image=tf_image,\n",
    "        pod_template_spec=pod_template_spec,\n",
    "    ).set_display_name(\"Visualization\")\n",
    "\n",
    "    prep_task = (\n",
    "        prep_op().after(prepare_tb_task\n",
    "                       ).set_display_name(\"Preprocess & Transform\")\n",
    "    )\n",
    "    confusion_matrix_url = f\"minio://{log_bucket}/{confusion_matrix_log_dir}\"\n",
    "    script_args = f\"model_name=resnet.pth,\" \\\n",
    "                  f\"confusion_matrix_url={confusion_matrix_url}\"\n",
    "    # For GPU, set number of gpus and accelerator type\n",
    "    ptl_args = f\"max_epochs=1, gpus=0, accelerator=None, profiler=pytorch\"\n",
    "    train_task = (\n",
    "        train_op(\n",
    "            input_data=prep_task.outputs[\"output_data\"],\n",
    "            cifar_script_args=script_args,\n",
    "            ptl_arguments=ptl_args\n",
    "        ).after(prep_task).set_display_name(\"Training\")\n",
    "    )\n",
    "    # For GPU uncomment below line and set GPU limit and node selector\n",
    "    # ).set_gpu_limit(1).add_node_selector_constraint\n",
    "    # ('cloud.google.com/gke-accelerator','nvidia-tesla-p4')\n",
    "\n",
    "    (\n",
    "        minio_op(\n",
    "            bucket_name=\"mlpipeline\",\n",
    "            folder_name=log_dir,\n",
    "            input_path=train_task.outputs[\"tensorboard_root\"],\n",
    "            filename=\"\",\n",
    "        ).after(train_task).set_display_name(\"Tensorboard Events Pusher\")\n",
    "    )\n",
    "\n",
    "    (\n",
    "        minio_op(\n",
    "            bucket_name=\"mlpipeline\",\n",
    "            folder_name=checkpoint_dir,\n",
    "            input_path=train_task.outputs[\"checkpoint_dir\"],\n",
    "            filename=\"\",\n",
    "        ).after(train_task).set_display_name(\"checkpoint_dir Pusher\")\n",
    "    )\n",
    "\n",
    "    minio_mar_upload = (\n",
    "        minio_op(\n",
    "            bucket_name=\"mlpipeline\",\n",
    "            folder_name=mar_path,\n",
    "            input_path=train_task.outputs[\"checkpoint_dir\"],\n",
    "            filename=\"cifar10_test.mar\",\n",
    "        ).after(train_task).set_display_name(\"Mar Pusher\")\n",
    "    )\n",
    "\n",
    "    (\n",
    "        minio_op(\n",
    "            bucket_name=\"mlpipeline\",\n",
    "            folder_name=config_prop_path,\n",
    "            input_path=train_task.outputs[\"checkpoint_dir\"],\n",
    "            filename=\"config.properties\",\n",
    "        ).after(train_task).set_display_name(\"Conifg Pusher\")\n",
    "    )\n",
    "\n",
    "    model_uri = str(model_uri)\n",
    "    # pylint: disable=unused-variable\n",
    "    isvc_yaml = \"\"\"\n",
    "    apiVersion: \"serving.kubeflow.org/v1beta1\"\n",
    "    kind: \"InferenceService\"\n",
    "    metadata:\n",
    "      name: {}\n",
    "      namespace: {}\n",
    "    spec:\n",
    "      predictor:\n",
    "        serviceAccountName: sa\n",
    "        pytorch:\n",
    "          storageUri: {}\n",
    "          resources:\n",
    "            requests: \n",
    "              cpu: 4\n",
    "              memory: 16Gi\n",
    "            limits:\n",
    "              cpu: 4\n",
    "              memory: 16Gi\n",
    "    \"\"\".format(\n",
    "        deploy, namespace, model_uri\n",
    "    )\n",
    "    \n",
    "    # For GPU inference use below yaml with gpu count and accelerator\n",
    "    gpu_count = \"1\"\n",
    "    accelerator = \"nvidia-tesla-p4\"\n",
    "    isvc_gpu_yaml = \"\"\"# pylint: disable=unused-variable\n",
    "    apiVersion: \"serving.kubeflow.org/v1beta1\"\n",
    "    kind: \"InferenceService\"\n",
    "    metadata:\n",
    "      name: {}\n",
    "      namespace: {}\n",
    "    spec:\n",
    "      predictor:\n",
    "        serviceAccountName: sa\n",
    "        pytorch:\n",
    "          storageUri: {}\n",
    "          resources:\n",
    "            requests: \n",
    "              cpu: 4\n",
    "              memory: 16Gi\n",
    "            limits:\n",
    "              cpu: 4\n",
    "              memory: 16Gi\n",
    "              nvidia.com/gpu: {}\n",
    "          nodeSelector:\n",
    "            cloud.google.com/gke-accelerator: {}\n",
    "\"\"\".format(deploy, namespace, model_uri, gpu_count, accelerator)\n",
    "    # Update inferenceservice_yaml for GPU inference\n",
    "    deploy_task = (\n",
    "        deploy_op(action=\"apply\", inferenceservice_yaml=isvc_yaml\n",
    "                 ).after(minio_mar_upload).set_display_name(\"Deployer\")\n",
    "    )\n",
    "    # Wait here for model to be loaded in torchserve for inference\n",
    "    sleep_task = sleep_op(5).after(deploy_task).set_display_name(\"Sleep\")\n",
    "    # Make Inference request\n",
    "    pred_task = (\n",
    "        pred_op(\n",
    "            host_name=isvc_name,\n",
    "            input_request=input_req,\n",
    "            cookie=cookie,\n",
    "            url=ingress_gateway,\n",
    "            model=model,\n",
    "            inference_type=\"predict\",\n",
    "        ).after(sleep_task).set_display_name(\"Prediction\")\n",
    "    )\n",
    "    (\n",
    "        pred_op(\n",
    "            host_name=isvc_name,\n",
    "            input_request=input_req,\n",
    "            cookie=cookie,\n",
    "            url=ingress_gateway,\n",
    "            model=model,\n",
    "            inference_type=\"explain\",\n",
    "        ).after(pred_task).set_display_name(\"Explanation\")\n",
    "    )\n",
    "\n",
    "    dsl.get_pipeline_conf().add_op_transformer(\n",
    "        use_k8s_secret(\n",
    "            secret_name=\"mlpipeline-minio-artifact\",\n",
    "            k8s_secret_key_to_env={\n",
    "                \"secretkey\": \"MINIO_SECRET_KEY\",\n",
    "                \"accesskey\": \"MINIO_ACCESS_KEY\",\n",
    "            },\n",
    "        )\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Compile the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "compiler.Compiler().compile(pytorch_cifar10, 'pytorch.tar.gz', type_check=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##  Execute the pipeline"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/html": [
       "<a href=\"http://istio-ingressgateway.istio-system.svc.cluster.local/pipeline/#/runs/details/4101b363-373c-4c29-9084-b5a77ff9de00\" target=\"_blank\" >Run details</a>."
      ],
      "text/plain": [
       "<IPython.core.display.HTML object>"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "run = client.run_pipeline(my_experiment.id, 'pytorch-cifar10', 'pytorch.tar.gz')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wait for inference service below to go to READY True state"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!kubectl get isvc $DEPLOY"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Get  the Inference service name"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 32,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'torchserve.kubeflow-user-example-com.example.com'"
      ]
     },
     "execution_count": 32,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "INFERENCE_SERVICE_LIST = ! kubectl get isvc {DEPLOY_NAME} -n {NAMESPACE} -o json | python3 -c \"import sys, json; print(json.load(sys.stdin)['status']['url'])\"| tr -d '\"' | cut -d \"/\" -f 3\n",
    "INFERENCE_SERVICE_NAME = INFERENCE_SERVICE_LIST[0]\n",
    "INFERENCE_SERVICE_NAME"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the deployed model for prediction request and save the output into a json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl -v -H \"Host: $INFERENCE_SERVICE_NAME\" -H \"Cookie: $COOKIE\" \"$INGRESS_GATEWAY/v1/models/$MODEL_NAME:predict\" -d @./cifar10/input.json > cifar10_prediction_output.json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! cat cifar10_prediction_output.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Use the deployed model for explain request and save the output into a json"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {},
   "outputs": [],
   "source": [
    "!curl -v -H \"Host: $INFERENCE_SERVICE_NAME\" -H \"Cookie: $COOKIE\" \"$INGRESS_GATEWAY/v1/models/$MODEL_NAME:explain\" -d @./cifar10/input.json > cifar10_explanation_output.json"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Model Interpretation using Captum Vis and Insights"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install dependencies for Captum Insights"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "!./install-dependencies.sh"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### import the necessary packages"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from PIL import Image\n",
    "import numpy as np \n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LinearSegmentedColormap\n",
    "import torchvision.transforms as transforms\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "import json\n",
    "import captum\n",
    "from captum.attr import LayerAttribution\n",
    "from captum.attr import visualization as viz\n",
    "import base64\n",
    "import os\n",
    "import io"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Read the prediction, explanation, and the class mapping file which saved during the prediction and expalain requests."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {},
   "outputs": [],
   "source": [
    "prediction_json = json.loads(open(\"./cifar10_prediction_output.json\", \"r\").read())\n",
    "\n",
    "explainations_json = json.loads(open(\"./cifar10_explanation_output.json\", \"r\").read())\n",
    "\n",
    "labels_path = './cifar10/class_mapping.json'\n",
    "with open(labels_path) as json_data:\n",
    "    idx_to_labels = json.load(json_data) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "count = 0\n",
    "for i in range(0, len(explainations_json[\"explanations\"])):\n",
    "\n",
    "    image = base64.b64decode(explainations_json[\"explanations\"][i][\"b64\"])       \n",
    "    fileName = 'captum_kitten_{}.jpeg'.format(count)\n",
    "\n",
    "    imagePath = ( os.getcwd() +\"/\" + fileName)\n",
    "    img = Image.open(io.BytesIO(image))\n",
    "    img = img.convert('RGB')\n",
    "    img.save(imagePath, 'jpeg', quality=100)\n",
    "    print(\"Saving \", imagePath)\n",
    "    count += 1\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from IPython.display import Image\n",
    "Image(filename='captum_kitten_0.jpeg') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename='captum_kitten_1.jpeg') "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "Image(filename='captum_kitten_2.jpeg') "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Captum Insights can also be used for visualization"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "##### Define the minio client for downloading the artifactes from minio storage ( model pth file and training file)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 41,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<minio.api.Minio at 0x7f45aa9f2d30>"
      ]
     },
     "execution_count": 41,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from minio import Minio\n",
    "from kubernetes import client, config\n",
    "import base64\n",
    "\n",
    "config.load_incluster_config()\n",
    "v1 = client.CoreV1Api()\n",
    "sec = v1.read_namespaced_secret(\"mlpipeline-minio-artifact\", NAMESPACE).data\n",
    "minio_accesskey = base64.b64decode(sec[\"accesskey\"]).decode('UTF-8')\n",
    "minio_secretkey = base64.b64decode(sec[\"secretkey\"]).decode('UTF-8')\n",
    "\n",
    "\n",
    "minio_config = {\n",
    "            \"HOST\": \"minio-service.kubeflow:9000\",\n",
    "            \"ACCESS_KEY\": minio_accesskey,\n",
    "            \"SECRET_KEY\": minio_secretkey,\n",
    "            \"BUCKET\": \"mlpipeline\",\n",
    "            \"FOLDER\": \"checkpoint_dir/cifar10\"}\n",
    "\n",
    "def _initiate_minio_client(minio_config):\n",
    "        minio_host = minio_config[\"HOST\"]\n",
    "        access_key = minio_config[\"ACCESS_KEY\"]\n",
    "        secret_key = minio_config[\"SECRET_KEY\"]\n",
    "        client = Minio(minio_host, access_key=access_key, secret_key=secret_key, secure=False)\n",
    "        return client\n",
    "\n",
    "client= _initiate_minio_client(minio_config)\n",
    "client"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 42,
   "metadata": {},
   "outputs": [],
   "source": [
    "def download_artifact_from_minio(folder: str, artifact: str):\n",
    "        artifact_name = artifact.split(\"/\")[-1]\n",
    "        result = client.fget_object(\n",
    "            minio_config[\"BUCKET\"],\n",
    "            os.path.join(folder, artifact_name),\n",
    "            artifact,\n",
    "        )"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "download_artifact_from_minio(minio_config[\"FOLDER\"],\"resnet.pth\")\n",
    "print(\"[INFO] Downloaded the Model Pth File.....\")\n",
    "download_artifact_from_minio(minio_config[\"FOLDER\"],\"cifar10_train.py\")\n",
    "print(\"[INFO] Downloaded the Model Classifier File.....\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Load the downloaded model pth file and classifer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "from cifar10_train import CIFAR10Classifier\n",
    "model = CIFAR10Classifier()\n",
    "\n",
    "model_pt_path =\"./resnet.pth\"\n",
    "model.load_state_dict(torch.load(model_pt_path,map_location=torch.device('cpu')))\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 39,
   "metadata": {},
   "outputs": [],
   "source": [
    "#Lets read two test images and make the prediction and use these images for captum Insights. \n",
    "\n",
    "imgs = ['./cifar10/kitten.png',\"./cifar10/horse.png\"]\n",
    "\n",
    "for img in imgs:\n",
    "    img = Image.open(img)\n",
    "    transformed_img = transform(img)\n",
    "    input_img = transform_normalize(transformed_img)\n",
    "    input_img = input_img.unsqueeze(0) # the model requires a dummy batch dimension\n",
    "\n",
    "    output = model(input_img)\n",
    "    output = F.softmax(output, dim=1)\n",
    "    prediction_score, pred_label_idx = torch.topk(output, 1)\n",
    "    pred_label_idx.squeeze_()\n",
    "    predicted_label = idx_to_labels[str(pred_label_idx.squeeze_().item())]\n",
    "    print('Predicted:', predicted_label, '/', pred_label_idx.item(), ' (', prediction_score.squeeze().item(), ')')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 48,
   "metadata": {},
   "outputs": [],
   "source": [
    "from captum.insights import AttributionVisualizer, Batch\n",
    "from captum.insights.attr_vis.features import ImageFeature\n",
    "\n",
    "# Baseline is all-zeros input - this may differ depending on your data\n",
    "def baseline_func(input):\n",
    "    return input * 0\n",
    "\n",
    "# merging our image transforms from above\n",
    "def full_img_transform(input):\n",
    "    i = Image.open(input)\n",
    "    i = transform(i)\n",
    "    i = transform_normalize(i)\n",
    "    i = i.unsqueeze(0)\n",
    "    i.requires_grad = True\n",
    "    return i\n",
    "\n",
    "\n",
    "input_imgs = torch.cat(list(map(lambda i: full_img_transform(i), imgs)), 0)\n",
    "\n",
    "visualizer = AttributionVisualizer(\n",
    "    models=[model],\n",
    "    score_func=lambda o: torch.nn.functional.softmax(o, 1),\n",
    "    classes=list(map(lambda k: idx_to_labels[k], idx_to_labels.keys())),\n",
    "    features=[\n",
    "        ImageFeature(\n",
    "            \"Photo\",\n",
    "            baseline_transforms=[baseline_func],\n",
    "            input_transforms=[],\n",
    "        )\n",
    "    ],\n",
    "    dataset=[Batch(input_imgs, labels=[3,7])]\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 50,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Fetch data and view Captum Insights at http://localhost:6080/\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "6080"
      ]
     },
     "execution_count": 50,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      " * Debugger is active!\n",
      " * Debugger PIN: 109-049-484\n",
      " * Running on http://127.0.0.1:6080/ (Press CTRL+C to quit)\n"
     ]
    }
   ],
   "source": [
    "visualizer.serve(debug=True,port=6080)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "scrolled": true
   },
   "source": [
    "#### Captum Insights output image "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![CIFAR10_Captum_Insights](./cifar10/CIFAR10_Captum_Insights.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Clean up\n",
    "### Delete Viewers, Inference Services and Completed pods"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "! kubectl delete --all isvc -n $NAMESPACE"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "! kubectl delete pod --field-selector=status.phase==Succeeded -n $NAMESPACE"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
